"""
回测引擎核心类
实现事件驱动的回测逻辑
"""
import os,sys
import time
from pandera.typing import DataFrame
from engine.object import BacktestResultSchema
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

import queue
from typing import Dict, List, Optional
from datetime import datetime
import pandas as pd
import numpy as np
from loguru import logger
from database.database import BaseDatabase
from engine.utility import symbolstr2list,listofstr2symbolstr
import copy
from .constant import OffsetType,AssetType,OrderType

if __name__ == "__main__":
    # 当作为脚本直接运行时，使用这些导入
    import sys
    from pathlib import Path
    
    # 将项目根目录添加到Python路径
    project_root = str(Path(__file__).parent.parent.absolute())
    if project_root not in sys.path:
        sys.path.insert(0, project_root)
    

    from engine.object import TradeData, CashFlowData, OrderData, BarData, OffsetType
    from engine.assets import Asset
    from engine.position import PortfolioPosition
    from engine.utility import listofstr2symbolstr,symbolstr2list
    from strategy.base import Strategy  
else:
    # 当作为模块导入时，使用相对导入
    from datetime import datetime

    from .object import TradeData, CashFlowData, OrderData, BarData, OffsetType
    from .assets import Asset
    from .position import PortfolioPosition
    from strategy.base import Strategy  


class BacktestEngine:
    """回测引擎核心类"""
    
    def __init__(
            self,
            start:str,
            end:str,
            offset:OffsetType = OffsetType.DAY,
            initial_capital: float = 1000000.0,
            commission_rate: float = 0.0003,
            slippage_rate: float = 0.0001,
            database:BaseDatabase = None,
            fill_vol_cap_pct:Optional[float] = 0.1,
            
        ):
        """
        初始化回测引擎
        
        Args:
            start: 回测起始日期,YYYYMMDD
            end: 回测结束日期,YYYYMMDD
            initial_capital: 初始资金
            commission_rate: 手续费率
            slippage_rate: 滑点率
            database: 数据库接口
            fill_vol_cap_pct: 每日最高可成交的数量占市场总成交数量比比例,None表示不限量
        """
        # 事件队列
        self.events: queue.Queue = queue.Queue()
        
        # 持仓管理
        self.initial_capital = initial_capital
        self.position = PortfolioPosition(initial_capital=initial_capital)
        
        # 交易成本参数
        self.commission_rate = commission_rate
        self.slippage_rate = slippage_rate
        
        # 回测数据
        self.market_data: Dict[str, Dict[datetime, BarData]] = {}
        self.assets: Dict[str, Asset] = {}
        self.factor_data: Dict[datetime, Dict[str, Dict[str, float]]] = {}
        # 回测结果
        self.current_date = self.current_time = None
        self.equity_curve: List[float] = []
        self.trades: List[TradeData] = []
        self.today_trades: List[TradeData] = []
        self.cashflows: List[CashFlowData] = []
        self.today_cashflows: List[CashFlowData] = []
        
        # 策略
        self.strategy: Optional[Strategy] = None
        self.active_orders: List[OrderData] = []
        self.database = database
        self.start=start
        self.end=end
        self.offset=offset
        
        self.fill_vol_cap_pct = fill_vol_cap_pct
        
        
        
        # 日志设置
        logger.add("logs/backtest.log", rotation="500 MB")
    
    def set_strategy(self, strategy: Strategy) -> None:
        """
        设置回测策略
        
        Args:
            strategy: 策略对象
        """
        # def add_strategy(self, strategy_class: type, setting: dict) -> None:
        #     """增加策略"""
        #     self.strategy_class = strategy_class
        #     self.strategy = strategy_class(
        #         self, strategy_class.__name__, copy(self.vt_symbols), setting
        #     )
        logger.info(f"加载策略: {strategy.name}")
        self.strategy = strategy
        self.strategy.set_engine(self)

    def load_data(self)->None:
        """
        加载回测数据
        """
        # 加载所有可用标的
        self.add_asset(None)
        self.add_market_data()
        # 加载因子数据
        all_symbols = ','.join(self.strategy.available_assets.keys())
        if all_symbols:
            # 使用引擎的数据库加载因子数据
            start_date = self.start.strftime("%Y%m%d") if isinstance(self.start, datetime) else self.start
            end_date = self.end.strftime("%Y%m%d") if isinstance(self.end, datetime) else self.end
            
            factor_data_list = self.database.load_factor_data(
                symbol=all_symbols,
                start=start_date,
                end=end_date
            )
            
            if factor_data_list:
                # 将因子数据处理成易于使用的格式并存储到引擎的factor_data中
                for factor_data in factor_data_list:
                    date = factor_data.timestamp
                    if date not in self.factor_data:
                        self.factor_data[date] = {}
                    
                    df = factor_data.factor_values
                    # 按照标的代码分组存储因子数据
                    for symbol in df.tradingcode.unique():
                        symbol_data = df[df.tradingcode == symbol]
                        if symbol not in self.factor_data[date]:
                            self.factor_data[date][symbol] = {}
                        
                        # 存储每个因子的值
                        for _, row in symbol_data.iterrows():
                            factor_name = row['factor_name']
                            factor_value = row['factor_value']
                            self.factor_data[date][symbol][factor_name] = factor_value
                
                logger.info(f"成功加载因子数据，时间范围: {start_date} - {end_date}")
            else:
                logger.warning("未能加载到因子数据")
    def initialize(self) -> None:
        """
        引擎初始化
        """
        logger.info('引擎初始化')
        self.strategy.initialize()  # 初始化策略

    
    def add_asset(self, asset: Asset | str = None) -> None:
        """
        添加资产
        
        Args:
            asset: 资产对象或资产代码,多个资产代码可以由逗号分隔
        """
        if asset is None:
            asset = self.database.get_all_asset_code() 
            logger.info(f"添加资产: {asset}")
        if isinstance(asset, str):
            # 如果输入是字符串（标的代码），使用数据库构建资产
            assets_list = symbolstr2list(asset) #去重
            for asset_ in assets_list:
                if asset_ in self.assets.keys():
                    assets_list.remove(asset_)
                    logger.info(f"资产 {asset_}已添加, 无需重复添加")
            if assets_list:
                asset = listofstr2symbolstr(assets_list)         
                assets = self.database.construct_asset(asset) 
            else:
                assets = []
            
            if assets:
                for asset_item in assets:
                    self.assets[asset_item.symbol] = asset_item
                    self.position.add_position(asset_item)
            else:
                logger.warning(f"无法从数据库加载资产: {asset}")

        else:
            # 如果输入是资产对象，直接添加
            if asset.symbol not in self.assets.keys():
                self.assets[asset.symbol] = asset
                self.position.add_position(asset)
    
    def add_market_data(self, symbol: str=None, 
                        data:Dict[datetime, BarData]=None) -> None:
        """
        添加市场数据
        
        Args:
            symbol: 证券代码
            data: 市场数据，可以为空, 如果不为空则优先加载传入数据
        """
        # 如果data不为None, 则优先加载传入的数据
        if data is not None:
            if isinstance(data, dict) and all(isinstance(k, datetime) for k in data.keys()):
                # 如果输入已经是{datetime: BarData}格式的字典
                if all(isinstance(v, BarData) for v in data.values()):
                    self.market_data[symbol] = data
                else:
                    raise TypeError("字典的值必须是BarData类型")
        else:
            if symbol is None and self.assets:
                #加载self.assets中的所有资产行情
                symbol = ','.join(self.assets.keys())
            else:
                raise Exception("请先添加资产")
            # 优先添加资产,重复的资产不会被添加
            self.add_asset(symbol)
            # 获取资产类型,多个标的这个版本资产类型必须保持一致
            assettype = self.assets[symbol.split(',')[0]].asset_type
            logger.info(f"加载K线数据: {symbol}, {self.start} - {self.end}")
            bars:List[BarData] = self.database.\
                                    load_bar_data(symbol, 
                                                  assettype, 
                                                  offset=self.offset,
                                                  start=self.start,
                                                  end=self.end)       
            # s_list = symbol.split(',')
            # symbol_dict = {k:v for k,v in zip(s_list, [{}]*len(s_list))}
            # for bar in bars:
            #     symbol_dict[bar.symbol].update({bar.timestamp:bar})
            symbol_dict = self._process_bar_data_to_market_dict(bars)
            self.market_data.update(symbol_dict)
            
    def _process_bar_data_to_market_dict(self, bars: List[BarData]) -> Dict[str, Dict[datetime, BarData]]:
        """
        将Bar数据列表处理成市场数据字典结构
        
        Args:
            bars: BarData对象列表
            
        Returns:
            Dict[str, Dict[datetime, BarData]]: 格式化后的市场数据字典
                - 外层键为交易标的代码
                - 内层键为时间戳
                - 值为对应的BarData对象
        """
        market_dict = {}
        
        # 按标的代码分组
        for bar in bars:
            if bar.symbol not in market_dict:
                market_dict[bar.symbol] = {}
            
            # 以时间戳为键存储BarData
            market_dict[bar.symbol][bar.timestamp] = bar
        
        return market_dict

    def place_order(self, 
                   symbol: str, 
                   quantity: float,
                   order_type: OrderType = OrderType.MARKET,
                   price: Optional[float] = None,
                   direction: int = 1,
                   asset_type: Optional[AssetType] = None):
        """
        下单
        
        参数:
            symbol: 证券代码
            quantity: 数量
            order_type: 订单类型，市价单或限价单
            price: 限价单价格
            direction: 方向，1为买入，-1为卖出
            asset_type: 资产类型
        """
        if symbol not in self.assets.keys():
            raise ValueError(f"未找到资产: {symbol}")
            
        # 创建订单
        order = OrderData(
                    timestamp=self.current_time,
                    asset_type=asset_type,
                    symbol=symbol,
                    order_type=order_type,
                    quantity=quantity,
                    price=price,
                    direction=direction,
        )
        


        self.active_orders.append(order)  
        logger.info(f"下单: {order}")   
    
    def execute_order(self, order: OrderData) -> None:
        """
        执行订单
        
        Args:
            order: 订单事件
        """
        # 获取当前市场数据
        symbol = order.symbol
        if symbol not in self.market_data:
            logger.warning(f"找不到资产 {symbol} 的市场数据")
            return
        
        if self.current_time not in self.market_data[symbol]:
            logger.warning(f"当前时间 {self.current_time} 没有 {symbol} 的市场数据")
            return
        
        # 获取当前价格
        current_bar = self.market_data[symbol][self.current_time]
        # 获取每日可成交量上限
        if self.fill_vol_cap_pct is not None:
            fill_vol_cap = self.market_data[symbol][self.current_time].volume * self.fill_vol_cap_pct
        else:
            fill_vol_cap = 10**20
        
        # 根据订单类型确定成交价格
        if order.order_type == OrderType.MARKET:
            # 市价单使用当前收盘价
            fill_price = current_bar.close_price
        elif order.order_type == OrderType.LIMIT:
            # 限价单需要检查是否可以成交
            if order.direction > 0:  # 买入
                if order.price < current_bar.close_price:
                    logger.info(f"限价单价格{order.price}低于市场价{current_bar.close_price}，无法成交")
                    return
            else:  # 卖出
                if order.price > current_bar.close_price:
                    logger.info(f"限价单价格{order.price}高于市场价{current_bar.close_price}，无法成交")
                    return
            fill_price = order.price
        else:
            logger.warning(f"不支持的订单类型: {order.order_type}")
            return
        

        
        # 创建成交事件
        if order.quantity > fill_vol_cap:
            logger.info(f"{order.symbol}: 订单数量 {order.quantity} 超过每日可成交量上限 {fill_vol_cap}")
        fillquantity = min(order.quantity,fill_vol_cap)

        # 计算交易成本
        commission = abs(fillquantity * fill_price) * self.commission_rate
        transaction_cost = abs(fillquantity * fill_price) * self.slippage_rate
        trade = TradeData(
            timestamp=self.current_time,
            asset_type=order.asset_type,
            symbol=order.symbol,
            quantity=fillquantity,
            direction=order.direction,
            price=fill_price,
            commission=commission,
            transaction_cost=transaction_cost,
            asset=self.assets[symbol]
        )

        
        # 更新持仓
        self.position.update_position_on_fill(trade)
        self.trades.append(trade)
        self.today_trades.append(trade)
        self.strategy.on_order_fill(trade) #成交事件
        # 记录日志
        logger.info(f"成交日期：{self.current_time}，订单执行: {trade.symbol}, 数量={trade.quantity}, " + \
                   f"方向={'买入' if trade.direction > 0 else '卖出'}, " + \
                   f"价格={fill_price:.4f}, 手续费={commission:.2f}, 交易成本={transaction_cost:.2f}")

    def update_market_data(self,date)->None:
        """
        更新持仓的相关数据为最新的市场数据
        """
        for symbol, data in self.market_data.items():
                        if date in data.keys():
                            self.position.update_market_data(
                                symbol, data[date].close_price, data[date].duration,data[date].YTM,
                                date
                            )
    def execute_all_orders(self) -> None:
        """执行所有订单"""
        while self.active_orders:
            self.execute_order(self.active_orders.pop(0))

                    
    def process_all_cashflows(self,cashflow_events:List[CashFlowData]):
        """
        处理所有现金流事件
        """
        if cashflow_events:
            logger.info(f"处理现金流事件：{cashflow_events}")
        # 优先处理现金流时间，因为现金流时间会影响可用金额
        for cashflow_event in cashflow_events:
            self.position.update_position_on_cashflow(cashflow_event)
  
    def run_backtest(self) -> DataFrame[BacktestResultSchema]:
        """
        运行回测
        
        Returns:
            pd.DataFrame: 回测结果
        """
        if self.strategy is None:
            raise ValueError("请先使用set_strategy设置回测策略")
        
        # 记录开始时间
        start_time = time.time()
        
        # 获取所有时间点
        all_dates = sorted(set().union(*[
            set(data.keys()) for data in self.market_data.values()
        ]))
        logger.info(f"回测时间范围: {all_dates[0]} 至 {all_dates[-1]}")
        results = []
        prev_value = None
        
        # 循环时间统计
        loop_times = []
        
        for date in all_dates:

            # 记录循环开始时间
            loop_start_time = time.time()
            
            self.today_trades = []
            self.today_cashflows = []
            self.current_time = date
            # 执行成交上一个bar的订单，会在这个交易日的市场价格下成交
            # 以今天的价格成交

            self.execute_all_orders()

            # 第一步：更新持仓的最新市场日期，和持仓价格
            self.update_market_data(date)
  
            
            # 第二步：处理每日现金流，更新当日账户头寸
            self.today_cashflows = self.position.check_cashflows(date)   

            self.process_all_cashflows(self.today_cashflows.copy())

                
            # 第三步：执行策略，根据昨日的行情价格和最新的可用资金，生成当日需要处理的订单
            for cashflow_event in self.today_cashflows:
                self.strategy.on_cashflow(cashflow_event) # 策略处理现金流事件

                
        
            if self.strategy.portfolio_flag:
                barset = {}
                for symbol,dict_bars in self.market_data.items():
                    bar = dict_bars.get(date)
                    if bar is not None:
                        bar_dict = {symbol:bar}
                        barset.update(bar_dict)
                if barset:
                    self.strategy.on_dataset(barset)


            else:
                for _,dict_bars in self.market_data.items():
                    bar =dict_bars.get(date)
                    if bar is not None:
                        self.strategy.on_data(bar)
                    


            
            

            # 记录每日净值,融资成本年化2%
            gpv = self.position.gross_total_value()
            npv = self.position.net_total_value()
            total_value =  npv - max(gpv-npv,0) * 0.02/365
                            
            self.equity_curve.append(total_value)
 
            # 计算收益率
            if prev_value is not None:
                returns = (total_value - prev_value) / prev_value
            else:
                returns = 0.0
            prev_value = total_value
            
            # 记录详细结果
            positions = self.position.get_position_sizes()
            weights = self.position.get_position_weights()
            duration = self.position.duration() #组合久期
            YTM = self.position.YTM() #组合YTM

            position_details = copy.deepcopy(self.position.get_position_details())
 
            accrued_interest = sum(
                pos.calculate_accrued_interest(date)
                for pos in self.position.positions.values() if pos.quantity != 0
            )
            
            # 添加当日结果
            results.append({
                'date': date,
                'total_value': total_value,
                'returns': returns,
                'cash': self.position.current_cash,
                'positions': positions,
                'weights': weights,
                'accrued_interest': accrued_interest,
                'trades': self.today_trades,
                'cashflows': self.today_cashflows,
                'position_detail': position_details,
                'duration': duration,
                'YTM': YTM
            })
            
            # 记录循环结束时间，并计算本次循环耗时
            loop_end_time = time.time()
            loop_duration = loop_end_time - loop_start_time
            loop_times.append(loop_duration)
        
        # 计算总耗时和平均循环耗时
        end_time = time.time()
        total_duration = end_time - start_time
        avg_loop_duration = sum(loop_times) / len(loop_times) if loop_times else 0
        
        # 输出时间统计信息
        logger.info(f"回测完成，总计耗时: {total_duration:.4f}秒")
        logger.info(f"共计处理 {len(all_dates)} 个交易日，平均每日耗时: {avg_loop_duration:.4f}秒")

        return pd.DataFrame(results)
    
    def reset_results(self, initial_capital: Optional[float] = None):
        """
        清空回测结果但保留市场数据，便于参数优化时重复使用已加载的数据
        
        Args:
            initial_capital: 可选，重置后的初始资金。如果不提供，则使用原初始资金
        
        Returns:
            None: 无返回值
        """
        # 清空事件队列
        # while not self.events.empty():
        #     self.events.get()
            
        # 重置持仓管理
        if initial_capital is None:
            initial_capital = self.position.initial_capital
        self.position = PortfolioPosition(initial_capital=initial_capital)
        # 添加资产
        for asset_item in self.assets.values():
            self.position.add_position(asset_item)
        # 重置回测结果相关变量
        self.current_time = None
        self.equity_curve = []
        self.trades = []
        self.today_trades = []
        self.cashflows = []
        self.today_cashflows = []
        self.active_orders = []
        self.strategy = None
        # 策略重新加载 

        # # 重新初始化策略（如果已设置）
        # if self.strategy:
        #     # 保存原策略的名称和参数
        #     strategy_name = self.strategy.name
        #     strategy_params = self.strategy.get_parameters()
        #     strategy_class = self.strategy.__class__
            
        #     # 重新创建策略实例
        #     self.strategy = strategy_class(name=strategy_name, **strategy_params)
        #     self.strategy.engine = self
            
        logger.info("已清空回测结果和策略，但是保留市场数据。")
        # return True
        
    def get_today_factor_value(self, factor_name=None):
        """
        获取当天所有标的的因子值
        
        Args:
            factor_name: 因子名称，如果不指定则返回所有因子值
            
        Returns:
            dict: 因子值字典，格式为 {symbol: {factor_name: factor_value}} 或 {symbol: factor_value}
        """
        if not hasattr(self, 'factor_data') or not self.factor_data:
            logger.warning("没有加载因子数据或因子数据为空")
            return {}
            
        # 如果当前时间不在因子数据中，返回空字典
        if self.current_time not in self.factor_data:
            logger.warning(f"当前时间 {self.current_time} 没有因子数据")
            return {}
            
        # 获取当天的因子数据
        today_factors = self.factor_data[self.current_time]
        
        # 如果指定了因子名称，只返回该因子的值
        if factor_name:
            result = {}
            for symbol, factors in today_factors.items():
                if factor_name in factors:
                    result[symbol] = factors[factor_name]
            return result
        else:
            # 返回所有因子值
            return today_factors
    def cover_cash_position(self)->None:
        """
        平头寸
        
        """
        self.position.current_cash
        pass
if __name__ == '__main__':
    from database.mysqldb import MysqlDatabase
    engine = BacktestEngine(start=20210101,end=20211231,database=MysqlDatabase())
    engine.add_asset('140024,160021')
    engine.add_market_data()
    # logger.info(engine.market_data)